{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset_io import *\n",
    "from color_channels import *\n",
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "import random as ra\n",
    "from tqdm import tqdm\n",
    "from profiler import *\n",
    "import random\n",
    "random.seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradient method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "nans=[np.NaN,np.NaN,np.NaN]\n",
    "\n",
    "def include_border(rect):\n",
    "    ((x1,y1),(x2,y2))=rect\n",
    "    border=IMG_SAFETY_BORDER\n",
    "    x1+=border\n",
    "    x2+=border\n",
    "    y1+=border\n",
    "    y2+=border\n",
    "    return ((x1,y1),(x2,y2))\n",
    "\n",
    "def add_border(image):\n",
    "    border=IMG_SAFETY_BORDER\n",
    "    return cv2.copyMakeBorder(image, border, border, border, border, cv2.BORDER_REFLECT)\n",
    "\n",
    "def update_optims(rect, img_height):\n",
    "    ((x1,y1),(x2,y2))=rect\n",
    "    w,h=(abs(x1-x2), abs(y1-y2))\n",
    "    global w_max,h_max,w_min,h_min,y1r_min,y2r_max\n",
    "    w_max=max(w,w_max)\n",
    "    h_max=max(h,h_max)\n",
    "    w_min=min(w,w_min)\n",
    "    h_min=min(h,h_min) \n",
    "    y1r=y1/img_height\n",
    "    y2r=y2/img_height\n",
    "    y1r_min=min(y1r, y1r_min)\n",
    "    y2r_max=max(y2r, y2r_max)\n",
    "    \n",
    "def draw_holes(img):\n",
    "    h,w,_=img.shape\n",
    "    hole_1=((x1,y1),(x2,y2))=((int(0.126*w), int(0.3056*h)),(int(0.409*w), int(h)))\n",
    "    img[y1:y2, x1:x2]=nans\n",
    "    hole_2=((x1,y1),(x2,y2))=((int(0.5906*w), int(0.3056*h)),(int(0.875*w), int(h)))\n",
    "    img[y1:y2, x1:x2]=nans\n",
    "\n",
    "def extract_pallet_rectangles_from_image(rects, imgs, filename, scene_name):\n",
    "    for j in range(len(rects)):\n",
    "        rect=rects[j]\n",
    "        rect=include_border(rect)\n",
    "        rect=correct_rect_ratio(rect)\n",
    "        ((x1,y1),(x2,y2))=add_margin(rect)\n",
    "        cropped_imgs = [img[y1:y2, x1:x2] for img in imgs]\n",
    "        cropped_imgs = [cv2.resize(img, (WIN_W,WIN_H)) for img in cropped_imgs]\n",
    "        save_image(cropped_imgs[0], scene_name, filename, \"pallet_rectangles_gradient\",str(j)+\"_ch1\",True)\n",
    "        save_image(cropped_imgs[1], scene_name, filename, \"pallet_rectangles_gradient\",str(j)+\"_ch2\",True)\n",
    "        global pallet_rectangles_gradient\n",
    "        pallet_rectangles_gradient[scene_name].append(cropped_imgs)\n",
    "            \n",
    "def extract_pallet_color_channels_from_image(rects, imgs, filename, scene_name, pallet_color):\n",
    "    colors=[pc for pc in pallet_color.split(',')]\n",
    "    for j in range(len(rects)):\n",
    "        rect=rects[j]\n",
    "        (x1,y1),(x2,y2)=rect\n",
    "        cropped_imgs = [img[y1:y2, x1:x2] for img in imgs]\n",
    "        [draw_holes(img) for img in cropped_imgs]\n",
    "        save_image(cropped_imgs[0], scene_name, filename, \"pallet_rectangles_color\",str(j)+\"_ch1\",True)\n",
    "        save_image(cropped_imgs[1], scene_name, filename, \"pallet_rectangles_color\",str(j)+\"_ch2\",True)\n",
    "        global pallets_color\n",
    "        flat_imgs = [np.reshape(img, (img.shape[0]*img.shape[1],img.shape[2])) for img in cropped_imgs]\n",
    "        fst_img=flat_imgs[0]\n",
    "        colorful_part=~(fst_img==0).all(axis=1)\n",
    "        flat_imgs = [img[colorful_part] for img in flat_imgs]\n",
    "        \n",
    "        color=colors[0] if len(colors)==1 else colors[j]\n",
    "        pallets_color[color][scene_name].append(flat_imgs)\n",
    "        \n",
    "def convert_to_gradient_channels(img):\n",
    "    img = cv2.UMat(img)\n",
    "    img = cv2.GaussianBlur(img, (3,3), sigmaX=0, sigmaY=0)\n",
    "    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY )\n",
    "    gray = cv2.equalizeHist(gray) #experimental\n",
    "    gradX = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3, scale=1 )\n",
    "    gradY = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3, scale=1 )\n",
    "    mag, ori = cv2.cartToPolar(gradX, gradY, angleInDegrees=True)\n",
    "\n",
    "    ori_ranges = [(0,60), (60,120), (120,180), (180,240), (240,300), (300,360)]\n",
    "    oris=[cv2.inRange(ori,l,r) for (l,r) in ori_ranges]\n",
    "    mags= [cv2.copyTo(mag, ori_x) for ori_x in oris]\n",
    "    channels=[gray,mag]+mags\n",
    "    channels=[cv2.UMat.get(ch) for ch in channels]\n",
    "    channels=np.moveaxis(channels, 0, -1)\n",
    "    return channels\n",
    "\n",
    "def convert_to_color_channels(img):\n",
    "    channels=to_color_channels(img, \"HSCMYb\")\n",
    "    return channels\n",
    "    \n",
    "def save_channels(channels, filename, scene_name, output_dir_name):\n",
    "    ch1,ch2=np.array_split(channels, 2,2)\n",
    "    save_image(ch1, scene_name, filename, output_dir_name,\"ch1\",True)\n",
    "    save_image(ch2, scene_name, filename, output_dir_name,\"ch2\",True)\n",
    "\n",
    "def extract_background_rectangles_from_image(rects, imgs, filename, scene_name):\n",
    "    if rects:\n",
    "        sample_pallet_w=0\n",
    "        y1s=[]\n",
    "        y2s=[]\n",
    "        hs=[]\n",
    "        for j in range(len(rects)):\n",
    "            rect=rects[j]\n",
    "            rect=correct_rect_ratio(rect)\n",
    "            ((x1,y1),(x2,y2))=add_margin(rect)\n",
    "            w,h=(abs(x1-x2), abs(y1-y2))\n",
    "            y1s.append(y1)\n",
    "            y2s.append(y2)\n",
    "            h=abs(y1-y2)\n",
    "            hs.append(h)\n",
    "        h=int(np.mean(hs))\n",
    "        w=5*h\n",
    "        y1=min(y1s)\n",
    "        y2=max(y2s)\n",
    "        \n",
    "        img_h,img_w,_=imgs[0].shape\n",
    "        sampling_x1=0\n",
    "        sampling_x2=img_w-w\n",
    "        sampling_upper_y1=0\n",
    "        sampling_upper_y2=y1-h\n",
    "        sampling_lower_y1=y2\n",
    "        sampling_lower_y2=img_h-h\n",
    "        \n",
    "        can_sample_upper=sampling_upper_y2>0\n",
    "        can_sample_lower=sampling_lower_y2>sampling_lower_y1\n",
    "        \n",
    "        samples_count=4\n",
    "        if (can_sample_upper or can_sample_lower) and sampling_x2>0: \n",
    "            for i in range(samples_count):\n",
    "                sampled_x=ra.randint(sampling_x1, sampling_x2)\n",
    "                sampling_done=False\n",
    "                if i%2 and can_sample_upper:\n",
    "                    sampled_y=ra.randint(sampling_upper_y1, sampling_upper_y2)\n",
    "                    sampling_done=True\n",
    "                elif can_sample_lower:\n",
    "                    sampled_y=ra.randint(sampling_lower_y1, sampling_lower_y2)\n",
    "                    sampling_done=True\n",
    "                if sampling_done:\n",
    "                    samples=[img[sampled_y:sampled_y+h, sampled_x:sampled_x+w] for img in imgs]\n",
    "                    samples=[cv2.resize(sample, (WIN_W,WIN_H)) for sample in samples]\n",
    "                    save_image(samples[0], scene_name, filename, \"backgrounds_gradient\",str(i)+\"_ch1\",True)\n",
    "                    save_image(samples[1], scene_name, filename, \"backgrounds_gradient\",str(i)+\"_ch2\",True)\n",
    "                    global background_rectangles_gradient\n",
    "                    background_rectangles_gradient[scene_name].append(samples)\n",
    "                   \n",
    "def natural(num):\n",
    "    return max(0, num)\n",
    "                \n",
    "def extract_background_color_channels_from_image(rects, imgs, filename, scene_name):\n",
    "    if rects:\n",
    "        y1s=[]\n",
    "        y2s=[]\n",
    "        x1s=[]\n",
    "        x2s=[]\n",
    "        for j in range(len(rects)):\n",
    "            rect=rects[j]\n",
    "            rect=correct_rect_ratio(rect)\n",
    "            ((x1,y1),(x2,y2))=add_margin(rect)\n",
    "            y1s.append(y1)\n",
    "            y2s.append(y2)\n",
    "            x1s.append(x1)\n",
    "            x2s.append(x2)\n",
    "        y1=natural(min(y1s))\n",
    "        y2=natural(max(y2s))\n",
    "        x1=natural(min(x1s))\n",
    "        x2=natural(max(x2s))\n",
    "        img_h,img_w,_=imgs[0].shape\n",
    "        \n",
    "        \n",
    "        upper_background=[img[0:y1, x1:x2] for img in imgs]\n",
    "        lower_background=[img[y2:img_h, x1:x2] for img in imgs]\n",
    "        backgrounds=[upper_background, lower_background]\n",
    "        global backgrounds_color\n",
    "        \n",
    "        for i in range(len(backgrounds)):\n",
    "            background=backgrounds[i]\n",
    "#             save_image(background[0], scene_name, filename, \"backgrounds_color\",str(i)+\"_ch1\",True)\n",
    "#             save_image(background[1], scene_name, filename, \"backgrounds_color\",str(i)+\"_ch2\",True)\n",
    "            background = [np.reshape(img, (img.shape[0]*img.shape[1],img.shape[2])) for img in background]\n",
    "            indices=range(len(background[0]))\n",
    "            subsample_indices=random.sample(indices, int(len(indices)/8)) #8 is manually adjusted to my RAM\n",
    "            background = [img[subsample_indices] for img in background]\n",
    "            backgrounds_color[scene_name].append(background)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Highest level functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image_optims_extraction(filename, scene_name, label_resolution, rects, pallet_color):\n",
    "    for rect in rects:\n",
    "        ((x1,y1),(x2,y2))=rect\n",
    "        rects=include_border(rect)\n",
    "        rects=correct_rect_ratio(rects)\n",
    "        rects=add_margin(rects)\n",
    "        img_height = label_resolution[1]\n",
    "        update_optims(rects, img_height)\n",
    "\n",
    "def process_image_gradient_channels_extraction(filename, scene_name, label_resolution, rects, pallet_color):\n",
    "    image=imread_resized(scene_name, filename, label_resolution)\n",
    "    channels=convert_to_gradient_channels(image)\n",
    "    save_channels(channels, filename, scene_name, \"channels_gradient\")\n",
    "    \n",
    "def process_image_color_channels_extraction(filename, scene_name, label_resolution, rects, pallet_color):\n",
    "    image=imread_resized(scene_name, filename, label_resolution)\n",
    "    channels=convert_to_color_channels(image)\n",
    "    save_channels(channels, filename, scene_name, \"channels_color\")\n",
    "    \n",
    "def process_image_gradient_extraction(filename, scene_name, label_resolution, rects, pallet_color):\n",
    "    channels=read_split_channels(scene_name, filename, \"channels_gradient\")\n",
    "    channels_bordered=[add_border(ch) for ch in channels]\n",
    "    extract_pallet_rectangles_from_image(rects, channels_bordered, filename, scene_name)\n",
    "    extract_background_rectangles_from_image(rects, channels, filename, scene_name)\n",
    "    \n",
    "def process_image_color_extraction(filename, scene_name, label_resolution, rects, pallet_color):\n",
    "    channels=read_split_channels(scene_name, filename, \"channels_color\")\n",
    "    extract_background_color_channels_from_image(rects, channels, filename, scene_name)\n",
    "    extract_pallet_color_channels_from_image(rects, channels, filename, scene_name, pallet_color)\n",
    "\n",
    "def process_data(fun):\n",
    "    [fun(*row) for row in tqdm(walk_dataset())]    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scenes_dict():\n",
    "    names=get_scene_names()\n",
    "    lists=[[] for n in names]\n",
    "    return dict(zip(names, lists))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "IMG_SAFETY_BORDER=130\n",
    "\n",
    "w_max=h_max=y2r_max=0\n",
    "w_min=h_min=y1r_min=99999\n",
    "\n",
    "pallets_color={color: get_scenes_dict() for color in COLORS}\n",
    "backgrounds_color=get_scenes_dict()\n",
    "background_rectangles_gradient=get_scenes_dict()\n",
    "pallet_rectangles_gradient=get_scenes_dict()\n",
    "\n",
    "MODE_GRADIENT_CHANNELS_EXTRACTION=process_image_gradient_channels_extraction\n",
    "MODE_COLOR_CHANNELS_EXTRACTION=process_image_color_channels_extraction\n",
    "MODE_GRADIENT_EXTRACTION=process_image_gradient_extraction\n",
    "MODE_COLOR_EXTRACTION=process_image_color_extraction\n",
    "MODE_OPTIMS=process_image_optims_extraction\n",
    "MODE=MODE_GRADIENT_EXTRACTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2092/2092 [01:42<00:00, 20.35it/s]\n"
     ]
    }
   ],
   "source": [
    "profiled('process_data(MODE)', globals(), locals())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "os.chdir(PROJECT_PATH) \n",
    "    \n",
    "for scene_name in get_scene_names():\n",
    "    if MODE==MODE_GRADIENT_EXTRACTION:\n",
    "        path = get_existing_clf_ds_filepath(\"background_rectangles_gradient\", scene_name)\n",
    "        np.save(path,background_rectangles_gradient[scene_name])\n",
    "        background_rectangles_gradient[scene_name]=None\n",
    "        \n",
    "        path = get_existing_clf_ds_filepath(\"pallet_rectangles_gradient\", scene_name)\n",
    "        np.save(path,pallet_rectangles_gradient[scene_name])\n",
    "        pallet_rectangles_gradient[scene_name]=None\n",
    "        \n",
    "    elif MODE==MODE_COLOR_EXTRACTION:\n",
    "        for color in COLORS:\n",
    "            path = get_existing_clf_ds_filepath(\"pallets_color\", scene_name, color)\n",
    "            np.save(path,pallets_color[color][scene_name])\n",
    "            pallets_color[color][scene_name]=None\n",
    "            \n",
    "        path = get_existing_clf_ds_filepath(\"backgrounds_color\", scene_name)\n",
    "        np.save(path,backgrounds_color[scene_name])\n",
    "        backgrounds_color[scene_name]=None\n",
    "        \n",
    "if MODE==MODE_OPTIMS:\n",
    "    print(\"w_max: \",w_max,\"h_max: \",h_max, \"y2r_max: \", y2r_max)\n",
    "    print(\"w_min: \",w_min,\"h_min: \",h_min, \"y1r_min: \", y1r_min)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
